import torch
import numpy as np
import os
import torch.utils.data as data
import cv2
class lowlight_loader(data.Dataset):

    def __init__(self,lowlight_image_path,ground_truth_path, low_ground_truth_path):
        self.size=512
        self.lowlight_image=lowlight_image_path
        self.ground_truth=ground_truth_path
        self.low_GT_path = low_ground_truth_path
        self.name=os.listdir(lowlight_image_path)
        self.name.sort(key=lambda x: int(x[:5]))
        print("Total training examples:", len(self.name))
    def __trans__(self,img):
        img = (np.asarray(img) / 255.0)
        img=torch.from_numpy(img).float()
        img=img.permute(2,0,1)
        return img
    def __getitem__(self, index):
        name = self.name[int(index)]
        prefix_name = name[:5]
        last_name = name[-4:]
        gt_name = prefix_name + last_name
        img = cv2.imread(os.path.join(self.lowlight_image, name))
        gt = cv2.imread(os.path.join(self.ground_truth, gt_name))
        low_gt = cv2.imread(os.path.join(self.low_GT_path, gt_name))
        return self.__trans__(img),self.__trans__(gt), self.__trans__(low_gt)
    def __len__(self):
        return len(self.name)


if __name__ == "__main__":
    lowlight_images_path = "/mnt/beta_2t/gys/RELLISUR/RELLISUR-Dataset/Train/LLLR"
    ground_truth_path = "/mnt/beta_2t/gys/RELLISUR/RELLISUR-Dataset/Train/NLHR/X2"
    low_ground_truth_path = "/mnt/beta_2t/gys/RELLISUR/RELLISUR-Dataset/Train/NLHR/X1"

    train_dataset = lowlight_loader(lowlight_images_path, ground_truth_path, low_ground_truth_path)
    img_lowlight, ground_truth, low_gt = train_dataset.__getitem__(0)
    print(img_lowlight.shape)
    print(ground_truth.shape)
    print(low_gt.shape)
    # name = "00001-2.5.png"
    # prefix_name = name[:5]
    # last_name = name[-4:]
    # print(last_name)
